# 数学计算工具包导入
import math

# torch以及torch.nn, torch.nn.functional
import torch
import torch.nn as nn
import torch.nn.functional as F

# torch中经典文本数据集有关的工具包
# 具体详情参考下方torchtext介绍
import torchtext

# torchtext中的数据处理工具, get_tokenizer用于英文分词
from torchtext.data.utils import get_tokenizer

# 已经构建完成的TransformerModel
from itcast_transformer import TransformerModel

# 将数据进行一个语料封装
TEXT = torchtext.data.Field(tokenize=get_tokenizer("basic_english"),
                            init_token='<sos>',
                            eos_token='<eos>',
                            lower=True)

# print(TEXT)

# 使用torchtext的数据集方法导入WikiText2数据
train_txt, val_txt, test_txt = torchtext.datasets.WikiText2.splits(TEXT)
# print(test_txt.examples[0].text[:10])

# 将训练集文本数据构建一个vocab对象, 可以使用vocab对象的stoi方法统计文本共包含的不重复的词汇总数
TEXT.build_vocab(train_txt)

# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 构建批次数据的函数
def batchify(data, batch_size):
    # data: 代表之前得到的文本数据(train_txt, val_txt, test_txt)
    # batch_size: 代表批次的样本量
    # 第一步使用TEXT的numericalize方法将单词映射成对应的连续数字
    data = TEXT.numericalize([data.examples[0].text])

    # 取得需要经过多少次的batch_size后能够遍历完所有的数据
    nbatch = data.size(0) // batch_size

    # 利用narrow方法对数据进行切割
    # 第1个参数代表横轴切割还是纵轴进行切割, 0代表横轴, 1代表纵轴
    # 第2个参数, 第3个参数分别代表切割的起始位置和终止位置
    data = data.narrow(0, 0, nbatch * batch_size)

    # 对data的形状进行转变
    data = data.view(batch_size, -1).t().contiguous()
    return data.to(device)

# x = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
# print(x.narrow(0, 0, 2))
# print(x.narrow(1, 1, 2))

# 设置训练数据的批次大小
batch_size = 20
# 设置验证数据和测试数据的批次大小
eval_batch_size = 10

# 获得训练数据，验证数据，测试数据
train_data = batchify(train_txt, batch_size)
val_data = batchify(val_txt, eval_batch_size)
test_data = batchify(test_txt, eval_batch_size)

# 设定句子的最大长度
bptt = 35

def get_batch(source, i):
    # source: 代表数据
    # i: 代表批次数

    # 确定句子的长度值
    seq_len = min(bptt, len(source) - 1 - i)

    # 首先得到源数据
    data = source[i:i+seq_len]
    # 然后得到目标数据
    target = source[i+1:i+1+seq_len].view(-1)
    return data, target

# source = test_data
# i = 1
# x, y = get_batch(source, i)
# print(x)
# print(y)

# 设置若干模型的超参数
# 通过TEXT.vocab.stoi方法获取不重复的词汇总数
ntokens = len(TEXT.vocab.stoi)

# 设置词嵌入维度的值等于200
emsize = 200

# 设置前馈全连接层的节点数等于200
nhid = 200

# 设置编码器层的层数等于2
nlayers = 2

# 设置多头注意力中的头数等于2
nhead = 2

# 设置置零的比率
dropout = 0.2

# 将参数传入TransformerModel中实例化模型
model = TransformerModel(ntokens, emsize, nhead, nhid, nlayers, dropout).to(device)

# 设定损失函数, 采用交叉熵损失函数
criterion = nn.CrossEntropyLoss()

# 设置学习率
lr = 5.0

# 设置优化器
optimizer = torch.optim.SGD(model.parameters(), lr=lr)

# 定义学习率调整器, 使用torch自带的lr_scheduler, 将优化器传入
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)

import time

# 构建训练函数
def train():
    # 首先开启训练模式
    model.train()
    # 定义初始损失值
    total_loss = 0
    # 设置打印间隔
    log_interval = 200
    # 获取当前开始时间
    start_time = time.time()
    # 遍历训练数据进行模型的训练
    for batch, i in enumerate(range(0, train_data.size(0) - 1, bptt)):
        # 通过前面的get_batch函数获取源数据和目标数据
        data, targets = get_batch(train_data, i)
        # 设置梯度归零
        optimizer.zero_grad()
        # 通过模型预测输出
        output = model(data)
        # 计算损失值
        loss = criterion(output.view(-1, ntokens), targets)
        # 进行反向传播
        loss.backward()
        # 先进行梯度规范化操作, 防止出现梯度爆炸或者梯度消失
        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        # 进行参数更新
        optimizer.step()
        # 将损失值进行累加
        total_loss += loss.item()
        # 打印日志信息
        if batch % log_interval == 0 and batch > 0:
            # 计算平均损失
            cur_loss = total_loss / log_interval
            # 计算训练到目前的耗时
            elapsed = time.time() - start_time
            # 打印日志信息
            print('| epoch {:3d} | {:5d}/{:5d} batches | '
                  'lr {:02.2f} | ms/batch {:5.2f} | '
                  'loss {:5.2f} | ppl {:8.2f}'.format(
                  epoch, batch, len(train_data) // bptt, scheduler.get_lr()[0],
                  elapsed * 1000 / log_interval,
                  cur_loss, math.exp(cur_loss)))

            # 每个打印批次结束后, 将总损失值清零
            total_loss = 0
            # 重新获取下一个打印轮次的开始时间
            start_time = time.time()

# 构建评估函数
def evaluate(eval_model, data_source):
    # eval_model: 代表每轮训练后产生的模型
    # data_source: 代表验证集数据或者测试机数据
    # 首先开启评估模式
    eval_model.eval()
    # 初始化总损失值
    total_loss = 0
    # 模型开启评估模式, 不进行反向传播求梯度
    with torch.no_grad():
        # 遍历验证数据
        for i in range(0, data_source.size(0) - 1, bptt):
            # 首先通过get_batch函数获取源数据和目标数据
            data, targets = get_batch(data_source, i)
            # 将源数据放入评估模型中, 进行预测
            output = eval_model(data)
            # 对输出张量进行变形
            output_flat = output.view(-1, ntokens)
            # 累加损失值
            total_loss += criterion(output_flat, targets).item()
    # 返回评估的总损失值
    return total_loss

# 首先初始化最佳模型的损失值为无穷大
best_val_loss = float("inf")

# 定义训练轮次
epochs = 3

# 定义最佳模型, 初始化为空
best_model = None

# 循环训练
for epoch in range(1, epochs + 1):
    # 获取当前轮次的开始时间
    start_time = time.time()
    # 直接调用训练函数进行模型的训练
    train()
    # 调用评估函数得到验证集的损失值
    val_loss = evaluate(model, val_data)
    # 打印信息
    print('-' * 90)
    print('| end of epoch {:3d} | time: {:5.2f}s | valid loss {:5.2f} | '.format(epoch, (time.time() - start_time), val_loss))
    print('-' * 90)
    # 通过比较当前轮次的损失值, 获取最佳模型
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        best_model = model
    # 每个轮次后调整优化器的学习率
    scheduler.step()


# 添加测试的流程代码
test_loss = evaluate(best_model, test_data)
print('-' * 90)
print('| End of training | test loss {:5.2f}'.format(test_loss))
print('-' * 90)
